import os
import cv2
import datetime as dt


class splitImg():

    def __init__(self, imgPath, labelsPath, savePath):
        self.rects = []
        self.classes = ['driver', 'assistant', 'person']
        self.imgPath = imgPath
        self.txtPath = labelsPath
        self.savePath = savePath

    def read_txt(self, filename, img):
        self.rects = []
        height, width = img.shape[:2]
        txt_path = os.path.join(self.txtPath, filename.replace('.jpg', '.txt').replace('.png', '.txt'))

        #读取txt文件获取信息
        with open(txt_path, 'r') as f:
            labels = f.read().splitlines()
        for i in labels:
            cl, x_centre, y_centre, w, h = i.split(' ')
            # 需要将数据类型转换成数字型
            cl, x_centre, y_centre, w, h = int(cl), float(x_centre), float(y_centre), float(w), float(h)
            name = self.classes[cl]  # 根据classes文件获取真实目标
            xmin = int(x_centre * width - w * width / 2)  # 坐标转换
            ymin = int(y_centre * height - h * height / 2)
            xmax = int(x_centre * width + w * width / 2)
            ymax = int(y_centre * height + h * height / 2)

            tmp = [name, xmin, ymin, xmax, ymax]  # 单个检测框
            self.rects.append(tmp)

    def selectROI(self, img):
        #截取roi区域
        for name, x1, y1, x2, y2 in self.rects:
            if name == "driver" or name == "assistant":
                roiImg = img[y1:y2, x1:x2]
                #根据系统时间命名图像
                now_time = dt.datetime.now().strftime("%Y%m%d%H%M%S%f")
                save_name = os.path.join(self.savePath + "/" + name + "/", str(now_time)[6:] + ".png")
                if not os.path.exists(self.savePath + "/" + name + "/"):
                    os.mkdir(self.savePath + "/" + name + "/")
                cv2.imwrite(save_name, roiImg)

    def run(self):
        #运行函数
        for file in os.listdir(self.imgPath):
            if file.split(".")[1] == "png" or file.split(".")[1] == "jpg":
                # 读取图片路径
                path = os.path.join(self.imgPath, file)
                img = cv2.imread(path)
                self.read_txt(file, img)
                self.selectROI(img)


if __name__ == '__main__':
    imgPath = "E:/pycharm/Codes/img/image"  #图片存储的根目录路径
    labelsPath = "E:/pycharm/Codes/img/label"  #xml存储的根目录路径
    savePath = "E:/pycharm/Codes/img/val_img"  #需要保存截取图像的位置
    task = splitImg(imgPath, labelsPath, savePath)
    task.run()
